import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler


# 加载数据
data = pd.read_csv("train.csv")


selected_features = [
    # "位置",
    "卧室数量",
    "卫的数量",
    "厅的数量",
    "总楼层",
    "房屋面积",
    "时间",
    "楼层",
    "Label",
]
data = data[selected_features]

print(data.isnull().sum())


X = data.drop("Label", axis=1)
y = data["Label"]
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)


X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=42
)

# 选择线性回归模型
model = LinearRegression()

# 训练模型
model.fit(X_train, y_train)

# 预测测试集结果
y_pred = model.predict(X_test)

# 计算均方误差
mse = mean_squared_error(y_test, y_pred)
print(f"模型的均方误差(MSE)为: {mse}")


# test
_data = pd.read_csv("test_noLabel.csv")

selected_features = [
    # "位置",
    "卧室数量",
    "卫的数量",
    "厅的数量",
    "总楼层",
    "房屋面积",
    "时间",
    "楼层",
]
p_data = _data[selected_features]
print(p_data.isnull().sum())


pX = scaler.fit_transform(p_data)

py = model.predict(pX)


_data["距离"] = py.tolist()

save_data = _data[["ID", "距离"]]

renamed_data = save_data.rename(columns={"ID": "ID", "距离": "Label"})

print(renamed_data.head())


renamed_data.to_csv("Answer_coshpr.csv", index=False)
